#!/usr/bin/env python
"""
通用定点巡航导航器，从 CSV 文件加载巡航路径，基于点云左右分布小范围避障
@author: CeasarSmj
@modify: ChatGPT
@date: 2025-4-19
"""

import os
import numpy as np
import open3d as o3d

import rospy

import sensor_msgs.point_cloud2 as pc2
from geometry_msgs.msg import Twist, Pose, Point, Quaternion, PoseWithCovarianceStamped
from sensor_msgs.msg import PointCloud2

import tf
import tf.transformations
import math
from std_msgs.msg import Header

# 导航部分与 ROS 解耦
class CruiseNavigator:
    # 导航参数
    position_tolerance = 0.35  # 米
    angular_tolerance = 0.15  # 弧度
    linear_speed = 0.4  # m/s
    rotation_speed_factor = 0.9  # 旋转速度比例
    avoidance_yaw = 0.1  # 避障时的小偏航角

    def __init__(self, key_points=None):
        self.current_idx = 0
        self.key_points = key_points or []

    def check_position_in_keypoints(self, position):
        """
        检查自身与当前及后续最多两个点的距离，如果在容差内则跳过，避免高速冲过造成误差
        """
        if self.current_idx >= len(self.key_points):
            return self.current_idx

        x, y, z, *_ = position
        max_offset = 0
        for i in (2, 1, 0):
            idx = self.current_idx + i
            if idx < len(self.key_points):
                tx, ty, tz = self.key_points[idx]
                dist = math.sqrt((x - tx) ** 2 + (y - ty) ** 2 + (z - tz) ** 2)
                if dist < self.position_tolerance:
                    max_offset = i + 1
                    break
        if max_offset:
            self.current_idx += max_offset
        return self.current_idx

    def check_cloud_in_cylinder(self, pc_data, start_pt, end_pt, diameter=0.1):
        """
        检查以 start_pt->end_pt 为轴, 直径 diameter 圆柱区域的点云分布
        返回: 'none','left','right','center'
        """
        return 'none'
        pts = np.asarray(pc_data.points)
        if pts.size == 0:
            return "none"
        p0 = np.array(start_pt)
        p1 = np.array(end_pt)
        axis = p1 - p0
        length = np.linalg.norm(axis)
        if length == 0:
            return "none"
        axis_dir = axis / length
        rel = pts - p0
        proj = np.dot(rel, axis_dir)
        mask_axial = (proj >= 0) & (proj <= length)
        radial2 = np.sum(rel**2, axis=1) - proj**2
        radius = diameter / 2.0
        mask = mask_axial & (radial2 <= radius**2)
        if not np.any(mask):
            return "none"
        pts_in = rel[mask]
        # 计算左右：在平面上，左侧向量为 [-dy, dx]
        perp = np.array([-axis_dir[1], axis_dir[0], 0.0])
        lateral = np.dot(pts_in, perp)
        n_left = np.sum(lateral > 0)
        n_right = np.sum(lateral < 0)
        if n_left and n_right:
            return "center"
        elif n_left > n_right:
            return "left"
        else:
            return "right"

    def update(self, position, pc_data):
        """
        返回: [linear_x, linear_y, angular_z]
        三种模式：1) 对齐航向  2) 等待（center）  3) 避障前进（left/right）
        """
        # 1) 切换关键点索引
        idx = self.check_position_in_keypoints(position)
        if idx >= len(self.key_points):
            return [0.0, 0.0, 0.0]
            
        rospy.loginfo("即将到达路径点: {}".format(self.key_points[self.current_idx]))

        x, y, z, _, _, yaw = position
        tx, ty, tz = self.key_points[idx]
        # 计算目标航向
        dx, dy = tx - x, ty - y
        desired_yaw = math.atan2(dy, dx)
        yaw_err = math.atan2(math.sin(desired_yaw - yaw), math.cos(desired_yaw - yaw))

        # 模式1：优先对齐航向
        if abs(yaw_err) > self.angular_tolerance:
            return [0.0, 0.0, self.rotation_speed_factor * yaw_err]

        # 模式2/3：局部点云避障
        mode = self.check_cloud_in_cylinder(
            pc_data, (x, y, z), (tx, ty, tz), diameter=0.6
        )
        if mode == "none":
            # 无障碍，直行
            return [self.linear_speed, 0.0, 0.0]
        if mode == "center":
            # 障碍居中，停止等待
            rospy.loginfo("barrier front, stop walking!")
            return [0.0, 0.0, 0.0]
        # 左/右避障：前进并小偏航
        yaw_delta = self.avoidance_yaw if mode == "left" else -self.avoidance_yaw
        return [self.linear_speed, 0.0, yaw_delta]


class ProcessorNode:
    log_file = "log.csv"
    path_file = "path.csv"

    def __init__(self):
        rospy.init_node("cruise_navigator_node")
        # self.publish_initial_pose([0.20137142916968612,0.2717263218759076,7.026585420050229],
        #                           [-0.0273024268890108,0.01295123576592705,3.1415926-1.6481527709571642])
        self.publish_initial_pose([0,0,0],
                                  [0,0,0])

        rospy.sleep(0.5)
        self.pc_data = None
        self.tf_listener = tf.TransformListener()

        base = os.path.dirname(os.path.abspath(__file__))
        self.log_path = os.path.join(base, self.log_file)
        self.path_path = os.path.join(base, self.path_file)

        pts = self.load_path_csv(self.path_path)
        self.navigator = CruiseNavigator(pts)

        rospy.Subscriber("/cur_scan_in_map", PointCloud2, self.pc_callback)
        self.pub = rospy.Publisher("/cmd_vel", Twist, queue_size=1)
        self.check_inv = False
        rospy.Timer(rospy.Duration(0.1), self.timer_callback)

        with open(self.log_path, "w") as f:
            f.write("timestamp,x,y,z,roll,pitch,yaw,linear_x,linear_y,angular_z\n")

    def publish_initial_pose(self, position, euler):
        pub = rospy.Publisher('/initialpose', PoseWithCovarianceStamped, queue_size=2)
        rospy.sleep(15.0)
        quat = tf.transformations.quaternion_from_euler(*euler)
        init = PoseWithCovarianceStamped()
        init.header = Header(stamp=rospy.Time.now(), frame_id='map')
        init.pose.pose.position = Point(*position)
        init.pose.pose.orientation = Quaternion(*quat)
        pub.publish(init)
        rospy.loginfo('---------------------------------------------------------------')        
        rospy.loginfo('---------------------------------------------------------------')
        rospy.loginfo('Initial pose published')
        rospy.loginfo('---------------------------------------------------------------')
        rospy.loginfo('---------------------------------------------------------------')
        self.ready_to_navigate = True

    def load_path_csv(self, csv_path):
        pts = []
        try:
            with open(csv_path, "r") as f:
                for line in f:
                    x, y, z = line.strip().split(",")
                    pts.append([float(x), float(y), float(z)])
        except Exception as e:
            rospy.logerr(f"读取路径文件失败: {e}")
        return pts

    def pc_callback(self, msg):
        self.pc_data = msg

    def tf_reader(self):
        try:
            trans, rot = self.tf_listener.lookupTransform("map", "body", rospy.Time(0))
            roll, pitch, yaw = tf.transformations.euler_from_quaternion(rot)
            return (*trans, roll, pitch, yaw)
        except:
            return None

    def pc_ros_to_o3d(self):
        pts = []
        for p in pc2.read_points(
            self.pc_data, field_names=("x", "y", "z"), skip_nans=True
        ):
            pts.append((p[0], p[1], p[2]))
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(np.array(pts))
        return pcd

    def timer_callback(self, event):
        if not self.pc_data:
            rospy.loginfo("No pointcloud data!!!!!!!!!!")
            return
        pos = self.tf_reader()
        if not pos:
            rospy.loginfo("No position data!!!!!!!!!!!!!!")
            return

        if not self.check_inv:
            rospy.loginfo(f"current position {pos},     \
                          first point {self.navigator.key_points[0]},   \
                          last point {self.navigator.key_points[-1]}")
            curr = np.array(pos[:3])
            start = np.array(self.navigator.key_points[0])
            end = np.array(self.navigator.key_points[-1])
            if np.linalg.norm(curr - end) < np.linalg.norm(curr - start):
                self.navigator.key_points.reverse()
            self.check_inv = True

        cloud_o3d = self.pc_ros_to_o3d()
        cmd = self.navigator.update(pos, cloud_o3d)

        rospy.loginfo(
            "Pos x=%.2f y=%.2f z=%.2f yaw=%.2f; Cmd lin=%.2f ang=%.2f",
            pos[0],
            pos[1],
            pos[2],
            pos[5],
            cmd[0],
            cmd[2],
        )
        with open(self.log_path, "a") as f:
            f.write(
                f"{rospy.Time.now()},{pos[0]},{pos[1]},{pos[2]},{pos[3]},{pos[4]},{pos[5]},{cmd[0]},{cmd[1]},{cmd[2]}\n"
            )

        twist = Twist()
        twist.linear.x, twist.linear.y, twist.angular.z = cmd
        self.pub.publish(twist)


if __name__ == "__main__":
    try:
        ProcessorNode()
        rospy.spin()
    except rospy.ROSInterruptException:
        pass